package cn.edu.sdjzu.knoverse.utils.rag;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.query.Query;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 增强型向量检索器，利用图数据库的关系信息来增强向量检索
 * 解决向量数据库的信息孤岛问题
 */
public class EnhancedVectorRetriever implements ContentRetriever {
    private static final Logger log = LoggerFactory.getLogger(EnhancedVectorRetriever.class);
    
    private final ContentRetriever graphRetriever;
    private final ContentRetriever vectorRetriever;
    private final ChatLanguageModel chatModel;

    // 用于增强查询的提示模板
    private final PromptTemplate enhanceQueryTemplate = PromptTemplate.from(
            "我将给你一个原始查询，以及从知识图谱中检索到的相关信息。" +
                    "请根据知识图谱的关系信息，重新构造一个更丰富的查询，以便从向量数据库中检索更相关的信息。\n\n" +
                    "原始查询: {{query}}\n\n" +
                    "知识图谱信息: {{graphInfo}}\n\n" +
                    "请生成一个增强后的查询，包含知识图谱中的关键实体、关系和属性，以便更精确地检索向量数据库。" +
                    "只返回增强后的查询文本，不要包含任何解释。"
    );

    public EnhancedVectorRetriever(
            ContentRetriever graphRetriever,
            ContentRetriever vectorRetriever,
            ChatLanguageModel chatModel
    ) {
        this.graphRetriever = graphRetriever;
        this.vectorRetriever = vectorRetriever;
        this.chatModel = chatModel;
    }

    @Override
    public List<Content> retrieve(Query query) {
        log.info("执行增强型检索: {}", query.text());

        // 1. 首先从图数据库检索相关信息
        List<Content> graphResults = graphRetriever.retrieve(query);
        if (graphResults.isEmpty()) {
            log.info("图数据库未返回结果，直接使用向量检索");
            return vectorRetriever.retrieve(query);
        }

        // 2. 从图数据库结果中提取文本信息
        String graphInfo = graphResults.stream()
                .map(content -> content.textSegment().text())
                .collect(Collectors.joining("\n"));

        log.info("图数据库检索结果: {}", graphInfo);

        // 3. 使用LLM增强查询
        Map<String, Object> variables = new HashMap<>();
        variables.put("query", query.text());
        variables.put("graphInfo", graphInfo);

        Prompt enhancedQueryPrompt = enhanceQueryTemplate.apply(variables);
        String enhancedQueryText = chatModel.chat(enhancedQueryPrompt.toUserMessage()).aiMessage().text();

        log.info("增强后的查询: {}", enhancedQueryText);

        // 4. 使用增强后的查询进行向量检索
        Query enhancedQuery = Query.from(enhancedQueryText);
        List<Content> vectorResults = vectorRetriever.retrieve(enhancedQuery);

        // 5. 合并结果，确保图数据库的结果也包含在内
        List<Content> combinedResults = new ArrayList<>(vectorResults);
        combinedResults.addAll(graphResults);

        return combinedResults;
    }
} 